import openpyxl
import requests
from bs4 import BeautifulSoup
from openpyxl.drawing.image import Image
import os


def download_image(image_url, save_path):
    response = requests.get(image_url, stream=True)
    if response.status_code == 200:
        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(1024):
                f.write(chunk)


def get_image_size(image_path):
    try:
        img = Image(image_path)
        return img.width, img.height
    except:
        return 0, 0


def calculate_scale_factor(img_width, img_height, cell_width, cell_height):
    img_width = img_width if img_width > 0 else cell_width
    img_height = img_height if img_height > 0 else cell_height
    width_scale = cell_width / img_width
    height_scale = cell_height / img_height
    return min(width_scale, height_scale)


def crawl_images_to_excel(url, excel_file, sheet_name):
    # Load the Excel workbook or create a new one if not exists
    try:
        wb = openpyxl.load_workbook(excel_file)
    except FileNotFoundError:
        wb = openpyxl.Workbook()

    # Select the sheet by name
    sheet = wb[sheet_name]

    # Send HTTP request and parse the HTML content
    response = requests.get(url)
    if response.status_code != 200:
        print(f"Failed to crawl images from {url}")
        return

    soup = BeautifulSoup(response.content, 'html.parser')

    # Find all image tags in the HTML content
    image_tags = soup.find_all('img')

    # Find all image titles (assuming they are in a specific tag, e.g., <h2>)
    title_tags = soup.find_all('h2')  # Replace 'h2' with the actual tag containing image titles

    # Create a directory to store downloaded images
    if not os.path.exists("images"):
        os.mkdir("images")

    # Insert images and titles into Excel
    for idx, img_tag in enumerate(image_tags):
        image_url = img_tag.get('src')
        title = title_tags[idx].text if idx < len(title_tags) else f"Image {idx + 1}"

        if not image_url:
            continue

        # Download the image to a local file
        image_filename = os.path.join("images", f"image_{idx}.jpg")
        download_image(image_url, image_filename)

        # Get the size of the image and the cell
        img_width, img_height = get_image_size(image_filename)
        cell_width = sheet.column_dimensions['B'].width  # Assuming the images are inserted in column B
        cell_height = sheet.row_dimensions[idx + 1].height  # Assuming one image per row

        # Calculate the scaling factor
        scale_factor = calculate_scale_factor(img_width, img_height, cell_width, cell_height)

        # Add the title and image to the worksheet
        cell_to_insert = f"A{idx + 1}"
        sheet[cell_to_insert] = title
        img = Image(image_filename)
        img.width = img_width * scale_factor
        img.height = img_height * scale_factor
        sheet.add_image(img, f"B{idx + 1}")

    # Save the workbook
    wb.save(excel_file)

    # Remove the downloaded image files (optional)
    # Comment out this section if you want to keep the images locally
    for idx in range(len(image_tags)):
        image_filename = os.path.join("images", f"image_{idx}.jpg")
        if os.path.exists(image_filename):
            os.remove(image_filename)


# Usage example:
website_url = "https://example.com"  # Replace with the target website URL
excel_file = "path/to/your/excel.xlsx"
sheet_name = "Sheet1"

crawl_images_to_excel(website_url, excel_file, sheet_name)
